package grpc_proxy_middleware

import (
	"encoding/json"
	"fmt"
	"log"

	"gin-gateway/core/utils"
	"gin-gateway/core/dao"
	"github.com/pkg/errors"
	"google.golang.org/grpc"
)

func GrpcJwtFlowCountMiddleware(serviceDetail *dao.ServiceDetail) func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		// 在处理器之间传递数据只能用 context，不能用 metadata.FromIncomingContext
		appInfos := ss.Context().Value("appInfo").(string)
		if len(appInfos) == 0 {
			if err := handler(srv, ss); err != nil {
				log.Printf("RPC failed with error %v\n", err)
				return err
			}
			return nil
		}

		appInfo := &dao.App{}
		if err := json.Unmarshal([]byte(appInfos), appInfo); err != nil {
			return err
		}
		RedisAppCounter, err := utils.FlowCountManager.GetRedisCounter(utils.FlowAppPrefix + appInfo.AppID)
		if err != nil {
			return err
		}
		RedisAppCounter.Increase()

		if appInfo.Qpd > 0 && RedisAppCounter.TotalCount > appInfo.Qpd {
			return errors.New(fmt.Sprintf("租户日请求量限流 limit:%v current:%v", appInfo.Qpd, RedisAppCounter.TotalCount))
		}
		if err = handler(srv, ss); err != nil {
			// log.Printf("RPC failed with error %v\n", err)
			return err
		}
		return nil
	}
}
